"""
多目标跟踪模块
使用ByteTrack算法进行人员跟踪
"""

import cv2
import numpy as np
import supervision as sv
import logging
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from collections import defaultdict

@dataclass
class TrackInfo:
    """跟踪信息数据类"""
    track_id: int
    bbox: Tuple[int, int, int, int]  # x1, y1, x2, y2
    center: Tuple[int, int]  # 中心点
    confidence: float
    class_id: int
    history: List[Tuple[int, int]]  # 历史轨迹点
    
class PersonTracker:
    """人员跟踪器"""
    
    def __init__(self, track_thresh: float = 0.25, 
                 track_buffer: int = 30,
                 match_thresh: float = 0.8,
                 frame_rate: int = 30):
        """
        初始化跟踪器
        
        Args:
            track_thresh: 跟踪置信度阈值
            track_buffer: 跟踪缓冲区大小
            match_thresh: 匹配阈值
            frame_rate: 帧率
        """
        self.track_thresh = track_thresh
        self.track_buffer = track_buffer
        self.match_thresh = match_thresh
        self.frame_rate = frame_rate
        
        # 初始化ByteTrack跟踪器
        self.tracker = sv.ByteTrack(
            track_activation_threshold=track_thresh,
            lost_track_buffer=track_buffer,
            minimum_matching_threshold=match_thresh,
            frame_rate=frame_rate
        )
        
        # 跟踪历史记录
        self.track_history: Dict[int, List[Tuple[int, int]]] = defaultdict(list)
        self.max_history_length = 50  # 最大历史长度
        
        # 日志
        self.logger = logging.getLogger(__name__)
        
    def update(self, detections: sv.Detections) -> sv.Detections:
        """
        更新跟踪器
        
        Args:
            detections: 检测结果
            
        Returns:
            带跟踪ID的检测结果
        """
        try:
            # 使用ByteTrack更新跟踪
            tracked_detections = self.tracker.update_with_detections(detections)
            
            # 更新跟踪历史
            self._update_track_history(tracked_detections)
            
            return tracked_detections
            
        except Exception as e:
            self.logger.error(f"跟踪更新失败: {e}")
            return detections
    
    def _update_track_history(self, detections: sv.Detections):
        """更新跟踪历史记录"""
        if detections.tracker_id is None:
            return
        
        for i, track_id in enumerate(detections.tracker_id):
            if track_id is None:
                continue
                
            # 计算中心点
            bbox = detections.xyxy[i]
            center_x = int((bbox[0] + bbox[2]) / 2)
            center_y = int((bbox[1] + bbox[3]) / 2)
            center = (center_x, center_y)
            
            # 添加到历史记录
            self.track_history[track_id].append(center)
            
            # 限制历史长度
            if len(self.track_history[track_id]) > self.max_history_length:
                self.track_history[track_id].pop(0)
    
    def get_track_info(self, detections: sv.Detections) -> List[TrackInfo]:
        """
        获取跟踪信息
        
        Args:
            detections: 带跟踪ID的检测结果
            
        Returns:
            跟踪信息列表
        """
        track_infos = []
        
        if detections.tracker_id is None:
            return track_infos
        
        for i, track_id in enumerate(detections.tracker_id):
            if track_id is None:
                continue
                
            bbox = detections.xyxy[i]
            center_x = int((bbox[0] + bbox[2]) / 2)
            center_y = int((bbox[1] + bbox[3]) / 2)
            
            track_info = TrackInfo(
                track_id=track_id,
                bbox=(int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])),
                center=(center_x, center_y),
                confidence=detections.confidence[i] if detections.confidence is not None else 0.0,
                class_id=detections.class_id[i] if detections.class_id is not None else 0,
                history=self.track_history[track_id].copy()
            )
            
            track_infos.append(track_info)
        
        return track_infos
    
    def get_track_history(self, track_id: int) -> List[Tuple[int, int]]:
        """获取指定跟踪ID的历史轨迹"""
        return self.track_history.get(track_id, [])
    
    def clear_old_tracks(self, current_track_ids: List[int]):
        """清理不活跃的跟踪记录"""
        active_ids = set(current_track_ids)
        inactive_ids = set(self.track_history.keys()) - active_ids
        
        for track_id in inactive_ids:
            del self.track_history[track_id]
    
    def reset(self):
        """重置跟踪器"""
        self.tracker = sv.ByteTrack(
            track_activation_threshold=self.track_thresh,
            lost_track_buffer=self.track_buffer,
            minimum_matching_threshold=self.match_thresh,
            frame_rate=self.frame_rate
        )
        self.track_history.clear()
        self.logger.info("跟踪器已重置")


class TrackingVisualizer:
    """跟踪结果可视化器"""
    
    def __init__(self):
        # 创建标注器
        self.box_annotator = sv.BoxAnnotator(
            thickness=2
        )
        
        self.label_annotator = sv.LabelAnnotator(
            text_scale=0.7,
            text_color=sv.Color.WHITE
        )
        
        self.trace_annotator = sv.TraceAnnotator(
            thickness=2,
            trace_length=30
        )
        
        # 颜色映射
        self.color_palette = sv.ColorPalette.DEFAULT
        
    def annotate_frame(self, frame: np.ndarray, detections: sv.Detections,
                      show_traces: bool = True) -> np.ndarray:
        """
        在帧上标注跟踪结果
        
        Args:
            frame: 输入帧
            detections: 带跟踪ID的检测结果
            show_traces: 是否显示轨迹
            
        Returns:
            标注后的帧
        """
        annotated_frame = frame.copy()
        
        # 绘制轨迹
        if show_traces and detections.tracker_id is not None:
            annotated_frame = self.trace_annotator.annotate(
                scene=annotated_frame,
                detections=detections
            )
        
        # 绘制边界框
        annotated_frame = self.box_annotator.annotate(
            scene=annotated_frame,
            detections=detections
        )
        
        # 创建标签
        labels = []
        if detections.tracker_id is not None:
            for i, track_id in enumerate(detections.tracker_id):
                confidence = detections.confidence[i] if detections.confidence is not None else 0.0
                if track_id is not None:
                    labels.append(f"ID:{track_id} {confidence:.2f}")
                else:
                    labels.append(f"Person {confidence:.2f}")
        else:
            labels = ["Person"] * len(detections)
        
        # 绘制标签
        annotated_frame = self.label_annotator.annotate(
            scene=annotated_frame,
            detections=detections,
            labels=labels
        )
        
        return annotated_frame
    
    def draw_track_info(self, frame: np.ndarray, track_infos: List[TrackInfo],
                       position: Tuple[int, int] = (10, 30)) -> np.ndarray:
        """
        绘制跟踪信息
        
        Args:
            frame: 输入帧
            track_infos: 跟踪信息列表
            position: 文本位置
            
        Returns:
            绘制信息后的帧
        """
        x, y = position
        
        # 绘制活跃跟踪数量
        active_tracks = len(track_infos)
        text = f"Active Tracks: {active_tracks}"
        
        cv2.putText(
            frame, text, (x, y),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2
        )
        
        # 绘制每个跟踪的ID
        for i, track_info in enumerate(track_infos[:5]):  # 最多显示5个
            track_text = f"ID:{track_info.track_id}"
            cv2.putText(
                frame, track_text, (x, y + 30 + i * 25),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2
            )
        
        return frame
    
    def draw_track_centers(self, frame: np.ndarray, track_infos: List[TrackInfo]) -> np.ndarray:
        """绘制跟踪中心点"""
        for track_info in track_infos:
            center = track_info.center
            cv2.circle(frame, center, 5, (0, 255, 0), -1)
            cv2.putText(
                frame, str(track_info.track_id), 
                (center[0] + 10, center[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2
            )
        
        return frame


if __name__ == "__main__":
    # 测试代码
    import time
    from detector import PersonDetector, DetectionVisualizer
    
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建检测器和跟踪器
    detector = PersonDetector()
    tracker = PersonTracker()
    detection_visualizer = DetectionVisualizer()
    tracking_visualizer = TrackingVisualizer()
    
    # 测试视频跟踪
    cap = cv2.VideoCapture(0)  # 使用摄像头测试
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 检测
        start_time = time.time()
        detections = detector.detect(frame)
        
        # 跟踪
        tracked_detections = tracker.update(detections)
        track_infos = tracker.get_track_info(tracked_detections)
        
        end_time = time.time()
        
        # 可视化
        annotated_frame = tracking_visualizer.annotate_frame(frame, tracked_detections)
        annotated_frame = tracking_visualizer.draw_track_info(annotated_frame, track_infos)
        
        # 显示FPS
        fps = 1 / (end_time - start_time)
        cv2.putText(
            annotated_frame, f"FPS: {fps:.1f}", (10, 150),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2
        )
        
        cv2.imshow("Person Tracking Test", annotated_frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows() 